from numpy import *


# data, handled
def loadDatas():
    postingList = [['my', 'dog', 'has', 'flea', 'problems', 'help', 'please'],
                   ['maybe', 'not', 'take', 'him', 'to', 'dog', 'park', 'stupid'],
                   ['my', 'dalmation', 'is', 'so', 'cute', 'I', 'love', 'him'],
                   ['stop', 'posting', 'stupid', 'worthless', 'garbage'],
                   ['mr', 'licks', 'ate', 'my', 'steak', 'how', 'to', 'stop', 'him'],
                   ['quit', 'buying', 'worthless', 'dog', 'food', 'stupid']]
    classVec = [0, 1, 0, 1, 0, 1]  # 1 is abusive, 0 not
    return postingList, classVec


def setsOfWordToVec(samples, data):
    vec = [0] * len(samples)
    for word in data:
        if word in samples:
            vec[samples.index(word)] = 1
    return vec


def getSamples(datas):
    samples = set([])
    for item in datas:
        samples = samples | set(item)
    return list(samples)


def train(mats, cls):
    trainNum = len(mats)
    sampleLen = len(mats[0])
    p1 = ones(sampleLen)  # has word count
    p0 = ones(sampleLen)
    p1s = 2  # all word count
    p0s = 2
    for i in range(trainNum):
        if cls[i] == 1:  # count
            p1 += mats[i]
            p1s = sum(mats[i])
        if cls[i] == 0:
            p0 += mats[i]
            p0s = sum(mats[i])
    p1v = log(p1 / p1s)
    p0v = log(p0 / p0s)
    pab = sum(cls) / len(cls)
    return p1v, p0v, pab


def classify(mat, p1v, p0v, pab):
    p1 = sum(mat * p1v) + pab
    p0 = sum(mat * p0v) + (1 - pab)
    if p1 > p0:
        return 1
    else:
        return 0


def test():
    datas, cls = loadDatas()
    samples = getSamples(datas)
    mat = []
    for item in datas:
        mat.append(setsOfWordToVec(samples, item))
    print mat
    p0v, p1v, pa = train(mat, cls)
    tests = ['my', 'love', 'good']
    testMat = setsOfWordToVec(samples, tests)
    print 'tes', classify(testMat, p0v, p1v, pa)


test()
